import time
from typing import Callable, List, Optional, Union

import numpy as np
from sutils.exp.logger import Logger


class DE:
    def __init__(
        self,
        eval_func: Callable[[np.ndarray], np.ndarray],
        n_dim: int,
        pop_size: int = 50,
        crl: float = 0.1,
        cru: float = 0.6,
        lower_bound: Union[np.ndarray, float] = -1.0,
        higher_bound: Union[np.ndarray, float] = 1.0,
        constraint_eq: List = [],
        constraint_ueq: List = [],
        target_fitness: float = np.inf,
    ):
        assert pop_size >= 3, "min pop_size is three"
        self._eval_func = eval_func
        self._n_dim = n_dim
        self._pop_size = pop_size
        self._crl = crl
        self._cru = cru
        if type(lower_bound) == float:
            self._lower_bound = np.array([lower_bound] * n_dim)
        else:
            self._lower_bound = lower_bound
        if type(higher_bound) == float:
            self._higher_bound = np.array([higher_bound] * n_dim)
        else:
            self._higher_bound = higher_bound

        self._have_constraint = len(constraint_eq) > 0 or len(constraint_ueq) > 0
        self._constraint_eq = constraint_eq
        self._constraint_ueq = constraint_ueq
        self._target_fitness = target_fitness
        self._pop = None

    def _create_population(self):
        # create the population
        population = np.random.uniform(
            low=self._lower_bound,
            high=self._higher_bound,
            size=(self._pop_size, self._n_dim),
        )
        return population

    def _get_fitness(self, pop: np.ndarray):
        y = self._eval_func(pop)

        if self._have_constraint:
            penalty_eq = np.array(
                [np.sum(np.abs([c_i(x) for c_i in self._constraint_eq])) for x in pop]
            )
            penalty_ueq = np.array(
                [
                    np.sum(np.abs([max(0, c_i(x)) for c_i in self._constraint_ueq]))
                    for x in pop
                ]
            )
            y -= 1e5 * penalty_eq + 1e5 * penalty_ueq
        return y

    def _mutation(self):
        '''
        V[i]=X[r1]+F(X[r2]-X[r3]),
        where i, r1, r2, r3 are randomly generated
        '''
        random_idx = []
        for _ in range(self._pop_size):
            idx = np.random.choice(self._pop_size, size=3, replace=False).tolist()
            idx.sort(reverse=True, key=lambda e: self._pop_fitness[e])
            random_idx.append(idx)
        random_idx = np.array(random_idx)
        r1, r2, r3 = random_idx[:, 0], random_idx[:, 1], random_idx[:, 2]
        # NOTE r1, r2, r3 should not be equal and f(r1) > f(r2) > f(r3)
        assert np.sum(r1 == r2) == 0
        assert np.sum(r1 == r3) == 0
        assert np.sum(r2 == r3) == 0

        if (
            np.sum(self._pop_fitness[r1] <= self._pop_fitness[r2]) == 0
            and np.sum(self._pop_fitness[r2] < self._pop_fitness[r3]) == 0
        ):
            factor = (
                0.1
                + 0.8
                * (self._pop_fitness[r2] - self._pop_fitness[r1])
                / (self._pop_fitness[r3] - self._pop_fitness[r1])
            ).reshape((-1, 1))
        else:
            factor = 0.5
        v_pop = self._pop[r1] + factor * (self._pop[r2] - self._pop[r3])

        # the lower & upper bound still works in mutation
        mask = np.random.uniform(
            low=self._lower_bound,
            high=self._higher_bound,
            size=(self._pop_size, self._n_dim),
        )
        v_pop = np.where(v_pop < self._lower_bound, mask, v_pop)
        v_pop = np.where(v_pop > self._higher_bound, mask, v_pop)
         
        return v_pop

    def _crossover(self, v_pop: np.ndarray):
        '''
        if rand < prob_crossover, use V, else use X
        '''
        fmin = self._pop_fitness.min()
        fmax = self._pop_fitness.max()
        fmean = self._pop_fitness.mean()
        index = self._pop_fitness > fmean
        prob_cr = np.zeros((self._pop_size,)) + self._crl
        if fmax - fmin > 1e-3:
            prob_cr[index] += (
                (self._cru - self._crl)
                * (self._pop_fitness[index] - fmin)
                / (fmax - fmin)
            )
        mask = np.random.rand(self._pop_size, self._n_dim) < prob_cr.reshape(-1, 1)
        u_pop = np.where(mask, v_pop, self._pop)
        return u_pop

    def _selection(self, u_pop: np.ndarray):
        '''
        greedy selection
        '''
        u_fitness = self._get_fitness(u_pop)
        index = self._pop_fitness > u_fitness
        pop = np.where(index.reshape(-1, 1), self._pop, u_pop)
        pop_fitness = np.where(index, self._pop_fitness, u_fitness)
        return pop, pop_fitness, u_fitness

    def run(
        self,
        max_iter: int = 200,
        record_history: bool = False,
        init_population: Optional[np.ndarray] = None,
    ):
        if record_history:
            best_fitness_history = []
            population_history = []
            fitness_history = []

        st_time = time.time()
        if init_population is not None:
            self._pop = init_population
        else:
            self._pop = self._create_population()
        self._pop_fitness = self._get_fitness(self._pop)
        Logger.info(
            f"init best_fitness: {self._pop_fitness.max()}, used time: {time.time() - st_time}"
        )
        Logger.debug(self._pop_fitness.__str__())

        if record_history:
            fitness_history.append(self._pop_fitness.copy())
            population_history.append(self._pop.copy())
            best_fitness_history.append(self._pop_fitness.max())

        index = 0
        while True:
            st_time = time.time()
            v_pop = self._mutation()
            u_pop = self._crossover(v_pop)
            self._pop, self._pop_fitness, u_fitness = self._selection(u_pop)

            if record_history:
                fitness_history.append(u_fitness.copy())
                population_history.append(u_pop.copy())
                best_fitness_history.append(self._pop_fitness.max())
            Logger.info(
                f"iter: {index}, best_fitness: {self._pop_fitness.max()}, used time: {time.time() - st_time}"
            )
            Logger.debug(u_fitness.__str__())
            if self._pop_fitness.max() > self._target_fitness:
                break
            index += 1
            if max_iter > 0 and index >= max_iter:
                break

        best_index = self._pop_fitness.argmax()
        best_x = self._pop[best_index]
        if self._have_constraint:
            best_y = self._eval_func(best_x.reshape(1, -1))[0]
        else:
            best_y = self._pop_fitness[best_index]

        if record_history:
            return (
                best_x,
                best_y,
                population_history,
                fitness_history,
                best_fitness_history,
            )
        else:
            return best_x, best_y
